package com.yhkj.data.datasource;

import com.yhkj.data.mq.message.DataSourceConfig;
import com.zaxxer.hikari.HikariConfig;
import com.zaxxer.hikari.HikariDataSource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.scheduling.annotation.Scheduled;

import javax.sql.DataSource;
import java.util.Iterator;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.locks.ReentrantReadWriteLock;

import static java.util.concurrent.TimeUnit.MINUTES;

/**
 * @author <a href="mailto:idler41@163.con">linfuxin</a> created on 2022-11-25 15:45:00
 */
@Slf4j
public class RoutingDataSource extends AbstractRoutingDataSource {

    private final Map<Object, HikariDataSourceHolder> targetDataSources = new ConcurrentHashMap<>();

    private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
    private final ReentrantReadWriteLock.ReadLock readLock = lock.readLock();
    private final ReentrantReadWriteLock.WriteLock writeLock = lock.writeLock();

    private final ConcurrentLinkedQueue<HikariDataSourceHolder> evictQueue = new ConcurrentLinkedQueue<>();

    @Value("${data-source.connection.expire:86400000}")
    private Long expireTime;

    public HikariDataSourceHolder buildDataSource(DataSourceConfig dataSourceConfig) {
        HikariConfig hikariConfig = new HikariConfig();
        hikariConfig.setJdbcUrl(dataSourceConfig.getJdbcUrl());
        hikariConfig.setUsername(dataSourceConfig.getUsername());
        hikariConfig.setPassword(dataSourceConfig.getPassword());
        hikariConfig.setDriverClassName(dataSourceConfig.getDriverClassName());
        hikariConfig.setConnectionTestQuery(dataSourceConfig.getValidationSql());
        hikariConfig.setMaximumPoolSize(dataSourceConfig.getPoolSize());
        hikariConfig.setMinimumIdle(1);
        hikariConfig.setIdleTimeout(MINUTES.toMillis(10));
        return new HikariDataSourceHolder(new HikariDataSource(hikariConfig));
    }

    @Override
    public void afterPropertiesSet() {
        // do nothing
    }

    @Override
    protected DataSource determineTargetDataSource() {
        Object lookupKey = determineCurrentLookupKey();
        HikariDataSourceHolder dsHolder = targetDataSources.get(lookupKey);
        if (dsHolder != null && dsHolder.getDataSource() != null) {
            return dsHolder.getDataSource();
        }
        throw new IllegalStateException("Cannot determine target DataSource for lookup key [" + lookupKey + "]");
    }

    @Override
    protected Object determineCurrentLookupKey() {
        DataSourceConfig dataSourceConfig = LocalRoutingKeyHolder.getContextKey();
        if (dataSourceConfig == null) {
            throw new RuntimeException("switch datasource key is null!");
        }
        String dsKey = getDataSourceKey(dataSourceConfig);
        try {
            readLock.lock();
            HikariDataSourceHolder dsHolder = targetDataSources.get(dsKey);
            if (dsHolder == null) {
                dsHolder = registerIfNecessary(dsKey, dataSourceConfig);
            }
            dsHolder.setAccessTime(System.currentTimeMillis());
            return dsKey;
        } finally {
            readLock.unlock();
        }
    }

    private HikariDataSourceHolder registerIfNecessary(String dsKey, DataSourceConfig dataSourceConfig) {
        return targetDataSources.computeIfAbsent(dsKey, k -> {
            if (log.isDebugEnabled()) {
                log.debug("register dataSource!key={}", dsKey);
            }
            return buildDataSource(dataSourceConfig);
        });
    }

    @Scheduled(cron = "${data-source.cron.evict:0 0 2 * * ?}")
    private void evictIfPossible() {
        try {
            writeLock.lock();
            // 标记出所有要销毁的连接池
            long now = System.currentTimeMillis();
            // expireTime时间内未访问
            for (Iterator<Map.Entry<Object, HikariDataSourceHolder>> it = targetDataSources.entrySet().iterator(); it.hasNext(); ) {
                Map.Entry<Object, HikariDataSourceHolder> entry = it.next();
                if (now - entry.getValue().getAccessTime() > expireTime) {
                    if (evictQueue.offer(entry.getValue())) {
                        it.remove();
                    }
                }
            }
        } finally {
            writeLock.unlock();
        }
        int evictCount = 0;
        HikariDataSourceHolder evictItem;
        while ((evictItem = evictQueue.poll()) != null) {
            try {
                evictItem.getDataSource().close();
                evictCount++;
            } catch (Exception e) {
                log.error("销毁数据源异常!", e);
            }
        }
        log.info("清理数据源个数={},剩余数据源={}", evictCount, targetDataSources.keySet().size());
    }

    private String getDataSourceKey(DataSourceConfig dataSourceConfig) {
        int index = dataSourceConfig.getJdbcUrl().indexOf("?");
        return index == -1 ? dataSourceConfig.getJdbcUrl() : dataSourceConfig.getJdbcUrl().substring(0, index);
    }
}
